package com.study.component.mybatis;

import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.util.JdbcConstants;
import java.util.List;
import java.util.Properties;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.mapping.SqlSource;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.DefaultReflectorFactory;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.factory.DefaultObjectFactory;
import org.apache.ibatis.reflection.wrapper.DefaultObjectWrapperFactory;


/**
 * @Author tanglingyu
 * @Date: 2021/07/01/ 15:28
 * @Description:基于druid的mybatis拦截器
 */
@Intercepts({@Signature(type = Executor.class, method = "update", args = {MappedStatement.class,
    Object.class})})
@Slf4j
public class MyBatisInterceptor implements Interceptor {

  private final MyBatisVisitor visitor = new MyBatisVisitor();


  @Override
  public Object intercept(Invocation invocation) throws Throwable {
    this.process(invocation);
    return invocation.proceed();
  }

  private void process(Invocation invocation) {
    Object[] args = invocation.getArgs();
    MappedStatement mappedStatement = (MappedStatement) args[0];
    SqlCommandType sqlCommandType = mappedStatement.getSqlCommandType();
    if (sqlCommandType != SqlCommandType.INSERT) {
      // 只处理insert
      return;
    }

    BoundSql boundSql = mappedStatement.getBoundSql(args[1]);
    String originalSql = boundSql.getSql();
    String modifiedSql = this.dispatch(originalSql);
    log.debug("修改SQL,修改后SQL:[{}], 修改后SQL:[{}]", originalSql, modifiedSql);

    resetSql2Invocation(invocation, modifiedSql);
  }


  private void resetSql2Invocation(Invocation invocation, String sql) {
    final Object[] args = invocation.getArgs();
    MappedStatement statement = (MappedStatement) args[0];
    Object parameterObject = args[1];
    BoundSql boundSql = statement.getBoundSql(parameterObject);
    MappedStatement newStatement = buildMappedStatement(statement, new BoundSqlSqlSource(boundSql));
    MetaObject msObject = MetaObject
        .forObject(newStatement, new DefaultObjectFactory(), new DefaultObjectWrapperFactory(),
            new DefaultReflectorFactory());
    msObject.setValue("sqlSource.boundSql.sql", sql);
    args[0] = newStatement;
  }

  private MappedStatement buildMappedStatement(MappedStatement ms, SqlSource newSqlSource) {
    MappedStatement.Builder builder =
        new MappedStatement.Builder(ms.getConfiguration(), ms.getId(), newSqlSource,
            ms.getSqlCommandType());
    builder.resource(ms.getResource());
    builder.fetchSize(ms.getFetchSize());
    builder.statementType(ms.getStatementType());
    builder.keyGenerator(ms.getKeyGenerator());
    if (ms.getKeyProperties() != null && ms.getKeyProperties().length != 0) {
      StringBuilder keyProperties = new StringBuilder();
      for (String keyProperty : ms.getKeyProperties()) {
        keyProperties.append(keyProperty).append(",");
      }
      keyProperties.delete(keyProperties.length() - 1, keyProperties.length());
      builder.keyProperty(keyProperties.toString());
    }
    builder.timeout(ms.getTimeout());
    builder.parameterMap(ms.getParameterMap());
    builder.resultMaps(ms.getResultMaps());
    builder.resultSetType(ms.getResultSetType());
    builder.cache(ms.getCache());
    builder.flushCacheRequired(ms.isFlushCacheRequired());
    builder.useCache(ms.isUseCache());

    return builder.build();
  }


  private String dispatch(String original) {
    // 处理sql，增加字段；
    List<SQLStatement> sqlStatements = SQLUtils.parseStatements(original, JdbcConstants.MYSQL);
    if (CollectionUtils.isNotEmpty(sqlStatements)) {
      for (SQLStatement sqlStatement : sqlStatements) {
        sqlStatement.accept(visitor);
      }
    }
    String newSql = SQLUtils.toSQLString(sqlStatements, JdbcConstants.MYSQL);
    return newSql;
  }


  @Override
  public Object plugin(Object target) {
    if (target instanceof Executor) {
      return Plugin.wrap(target, this);
    }
    return target;
  }

  @Override
  public void setProperties(Properties properties) {

  }


  class BoundSqlSqlSource implements SqlSource {

    private BoundSql boundSql;

    public BoundSqlSqlSource(BoundSql boundSql) {
      this.boundSql = boundSql;
    }

    @Override
    public BoundSql getBoundSql(Object parameterObject) {
      return boundSql;
    }
  }

}

